import os, requests, json, csv
import pandas as pd
from typing import Tuple, List
from llm_engine.config import API_BASE, API_KEY
from zai import ZhipuAiClient


def load_excel_data(excel_file: str) -> Tuple[List[str], List[str]]:
    
    df_companies = pd.read_excel(excel_file, sheet_name="清单")
    company_list = df_companies.iloc[0:, 1].dropna().astype(str).tolist()
    
    df_types = pd.read_excel(excel_file, sheet_name="类型")
    label_list = df_types.iloc[0:, 0].dropna().astype(str).tolist()
    
    return company_list, label_list


def classify_company(company_name: str, candidate_labels) -> dict:
    client = ZhipuAiClient(api_key=API_KEY)
    
    # 先调用网络搜索和总结
    tools = [{
        "type": "web_search",
        "web_search": {
            "enable": True,
            "search_engine": "search_std",
            "search_result": True,
            "search_prompt": "你是一位企业性质分析专家。请简洁朴素的总结网络搜索{search_result}中的关键信息。返回内容100字以内, 不要使用markdown",
            "count": 5,
            "search_recency_filter": "noLimit",
            "content_size": "low"
        }
    }]
    messages = [
        {"role": "user","content": f"{company_name}的主要产品和服务是什么? 用100字左右进行简洁描述"},
    ]
    
    response = client.chat.completions.create(
        model="glm-4.5-air",  # 模型标识符
        messages=messages,  # 用户消息
        tools=tools,         # 工具参数
        thinking={
            "type": "disabled",    # 启用深度思考模式
        },
    )
    
    choice = response.choices[0] 
    info = choice.message.content 

    # print(info)
        
    
    system_prompt = f"""
        你是一个企业分类助手。  
        你需要根据公司介绍，将该公司分类到下列选项中最恰当的一个：  
        {candidate_labels}

        要求：    
        1. 仅从以上枚举集合中选择一个最合适的, 如果没有合适的则选择"其他", 不要返回任何枚举范围之外的分类结果。
        2. 严格限制返回本文, 仅返回分类结果和可信度(0-100, 越高越可信), 用|符号分割, 不要返回任何其他内容或者符号. 结果示例: 网络产品|80  
    """
    
    user_prompt = f"请根据介绍{info}, 为以下公司分类：{company_name}"
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    
    
    
    # 调用API获取响应
    response = client.chat.completions.create(
        model="glm-4.5-air",  # 模型标识符
        messages=messages,  # 用户消息
        thinking={
            "type": "disabled",    # 启用深度思考模式
        },
    )
    choice = response.choices[0] 
    content = choice.message.content 
    
    label = content.split("|")[0].replace(" ", "").replace("\n", "").replace("\r", "")
    trsutable = content.split("|")[1].replace(" ", "").replace("\n", "").replace("\r", "")
    avaliable = "是" if label in candidate_labels else "否"
    result = {
        "company": company_name, 
        "label": label,
        "trsutable":trsutable,
        "avaliable":avaliable,
        "description":info.replace(" ", "").replace("\n", "").replace("\r", "").replace(",", "，")
    }
    return result

if __name__ == '__main__':
    current_dir = os.path.dirname(os.path.abspath(__file__))
    excel_path = os.path.join(current_dir, "20250902.xlsx")
    output_csv = os.path.join(current_dir, "company_labeled.csv")
    
    
    companies, labels = load_excel_data(excel_path)
    cnt = 270
    with open(output_csv, "a", newline="", encoding="utf-8-sig") as f:
        writer = csv.writer(f)
        writer.writerow(["company", "label", "avaliable", "trsutable","description"])
        for company in companies[270:]:
            try:
                result = classify_company(company, labels)
            except Exception as e:
                result = {"company": company, "label": "FIXME大模型报错", "avaliable": "否", "trsutable": "0", "description": str(e)}
                print(e)
            print(result)
            print(f"{cnt} / {len(companies)}")
            cnt += 1
            writer.writerow([result["company"], result["label"], result["avaliable"], result["trsutable"], result["description"]])
            